# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import base64
import datetime
import enum
import hashlib
import hmac
import json
import logging
import time
import typing
import urllib.parse

import requests
import requests.auth

from coin.exchange.base.kr_rest.auth_base import AuthKey
from coin.exchange.base.kr_rest.rest_client_base import RestUpdate
from coin.exchange.util.async_http_client import create_async_http_client
from coin.exchange.util.rate_limit import RateLimiter


class HuobiOrderState(enum.Enum):
  PRE_SUBMITTED = 'pre-submitted'
  SUBMITTING = 'submitting'
  SUBMITTED = 'submitted'
  PARTIAL_FILLED = 'partial-filled'
  FILLED = 'filled'
  CANCELLING = 'cancelling'
  PARTIAL_CANCELED = 'partial-canceled'
  CANCELED = 'canceled'


class HuobiOrderType(enum.Enum):
  BUY_MARKET = 'buy-market'
  SELL_MARKET = 'sell-market'
  BUY_LIMIT = 'buy-limit'
  SELL_LIMIT = 'sell-limit'
  BUY_IOC = 'buy-ioc'
  SELL_IOC = 'sell-ioc'
  BUY_LIMIT_MAKER = 'buy-limit-maker'
  SELL_LIMIT_MAKER = 'sell-limit-maker'


class HuobiAccountType(enum.Enum):
  SPOT = 'spot'
  POINT = 'point'
  FUTURES = 'futures'
  MARGIN = 'margin'
  SUPER_MARGIN = 'super-margin'


HUOBI_API_HOST = 'api.huobi.pro'


class HuobiAuth(requests.auth.AuthBase):
  def __init__(self, key_file):
    self._key = AuthKey.from_file(key_file)

  def _get_signed_params(self, api_host, method, endpoint, params=None):
    assert method in ('GET', 'POST'), method
    timestamp = datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%S')
    if api_host == 'xiguaapi.hbdm.com':
      api_host = 'api.hbdm.com'
    full_params = {
        'AccessKeyId': self._key.access_key,
        'SignatureMethod': 'HmacSHA256',
        'SignatureVersion': '2',
        'Timestamp': timestamp
    }
    if method == 'GET':
      full_params.update(params)

    sorted_params = sorted(full_params.items(), key=lambda d: d[0])
    encode_params = urllib.parse.urlencode(sorted_params)
    payload = [method, api_host, endpoint, encode_params]
    payload = '\n'.join(payload).encode()
    signature = hmac.new(self._key.secret_key_bytes, payload, digestmod=hashlib.sha256)
    signature = base64.b64encode(signature.digest()).decode()
    full_params['Signature'] = signature
    return full_params

  def __call__(self, req: requests.PreparedRequest):
    assert req.method in ('GET', 'POST'), req.method
    method = req.method
    parsed_result = urllib.parse.urlparse(req.url)
    api_host = parsed_result.netloc
    path = parsed_result.path
    # noinspection PyArgumentList
    url = urllib.parse.ParseResult(parsed_result.scheme,
                                   parsed_result.netloc,
                                   parsed_result.path,
                                   '',
                                   '',
                                   '').geturl()

    if method == 'GET':
      params = urllib.parse.parse_qsl(parsed_result.query)
      params = self._get_signed_params(api_host, req.method, path, params)
      req.prepare_url(url, params)
    else:
      params = self._get_signed_params(api_host, req.method, path)
    req.prepare_url(url, params)
    return req


class HuobiSyncQueryer(object):
  def __init__(self, key_file, *, api_host=None, disable_ssl=False, timeout=10):
    self._timeout = timeout
    self._auth = HuobiAuth(key_file)
    scheme = 'http://' if disable_ssl else 'https://'
    self._api_host = scheme + (api_host or HUOBI_API_HOST)
    self.rate_limit = RateLimiter(100, 10)
    self._logger = logging.getLogger(__name__)

  def auth_request(self, method, path, params):
    assert method in ('GET', 'POST'), method
    data = None
    if method == 'POST':
      data = params
      params = None

    url = urllib.parse.urljoin(self._api_host, path)
    response = None
    try:
      response = requests.request(method,
                                  url,
                                  params=params,
                                  json=data,
                                  auth=self._auth,
                                  timeout=self._timeout)
      content = response.json()
      return content
    except requests.exceptions.RequestException:
      self._logger.exception('RequestException: %s')
      raise
    except json.JSONDecodeError:
      self._logger.exception('Json decode: %s', response.content)
      raise
    except Exception as e:
      self._logger.exception('Unknown exception, %s', e)
      raise

  def _query_nocheck(self, method, path, params):
    req_timestamp = time.time()
    content = self.auth_request(method, path, params)
    res_timestamp = time.time()
    status = requests.codes.ok
    update = RestUpdate(req_timestamp, res_timestamp, content, status)
    return update

  def query(self, method, path, params=None, disable_rate_limit=False):
    if not disable_rate_limit:
      while not self.rate_limit.check():
        time.sleep(0.1)
    else:
      self.rate_limit.check()
    return self._query_nocheck(method, path, params)


class HuobiAsyncQueryer(object):
  def __init__(
      self,
      key_file,
      *,
      api_host=None,
      disable_ssl=False,
      timeout=10,
      max_connections=4,
      rest_lib='pycurl',
  ):
    self._timeout = timeout
    self._http_client = create_async_http_client(rest_lib=rest_lib, max_connections=max_connections)
    self._auth = HuobiAuth(key_file)
    scheme = 'http://' if disable_ssl else 'https://'
    self._api_host = scheme + (api_host or HUOBI_API_HOST)

  def query(self, method, path, params=None):
    assert method in ('GET', 'POST'), method
    data = None
    if method == 'POST':
      data = params
      params = None

    url = urllib.parse.urljoin(self._api_host, path)
    future = self._http_client.fetch(method=method,
                                     url=url,
                                     params=params,
                                     json=data,
                                     auth=self._auth,
                                     timeout=self._timeout)
    return future


class HuobiNativePrivateClient(object):
  def __init__(self,
               *,
               key_file,
               api_host=None,
               disable_ssl=False,
               timeout=10,
               max_connections=8,
               use_async=False,
               rest_lib='pycurl'):
    rest_lib = rest_lib or None
    if use_async:
      self._auth_queryer = HuobiAsyncQueryer(key_file,
                                             api_host=api_host,
                                             disable_ssl=disable_ssl,
                                             timeout=timeout,
                                             max_connections=max_connections,
                                             rest_lib=rest_lib)
    else:
      self._auth_queryer = HuobiSyncQueryer(key_file,
                                            api_host=api_host,
                                            disable_ssl=disable_ssl,
                                            timeout=timeout)

  def _query(self, method, path, params=None):
    return self._auth_queryer.query(method, path, params=params)

  def query_accounts(self):
    method = 'GET'
    url = '/v1/account/accounts'
    update = self._query(method, url)
    return update

  def query_account_balance(self, account_id):
    method = 'GET'
    url = '/v1/account/accounts/%s/balance' % account_id
    update = self._query(method, url)
    return update

  def query_deposit_withdraw(self, currency, type):
    assert type in ('deposit', 'withdraw'), type
    method = 'GET'
    url = '/v1/query/deposit-withdraw'
    params = {'currency': currency, 'type': type}
    update = self._query(method, url, params=params)
    return update

  def query_orders(self,
                   symbol,
                   states: typing.List[HuobiOrderState],
                   huobi_order_types: typing.List[HuobiOrderType],
                   start_date=None,
                   end_date=None,
                   from_id=None,
                   direct=None,
                   size=None):
    method = 'GET'
    url = '/v1/order/orders'
    assert direct in (None, 'prev', 'next'), direct

    states_str = ','.join(s.value for s in states)
    order_types_str = ','.join(t.value for t in huobi_order_types)
    params = {
        "symbol": symbol,
        "states": states_str,
        "types": order_types_str,
        "start_date": start_date,
        "end_date": end_date,
        "from": from_id,
        "direct": direct,
        "size": size
    }
    update = self._query(method, url, params)
    return update

  def query_single_order(self, order_id):
    method = 'GET'
    url = '/v1/order/orders/{order_id}'.format(order_id=order_id)
    update = self._query(method, url)
    return update

  def query_open_orders(
      self,
      account_id=None,
      symbol=None,
      side=None,
      size=None,
  ):
    method = 'GET'
    url = '/v1/order/openOrders'
    assert bool(account_id) == bool(symbol), \
        'account-id and symbol must be both None or both set'
    if size is not None:
      assert 0 <= size <= 500, size
    side = side if side is None else side.value
    params = {
        'account-id': account_id,
        'symbol': symbol,
        'side': side,
        'size': size,
    }
    update = self._query(method, url, params)
    return update

  def query_order_match_results(self,
                                symbol=None,
                                order_types=None,
                                start_date=None,
                                end_date=None,
                                from_id=None,
                                direct=None,
                                size=None):
    method = 'GET'
    url = '/v1/order/matchresults'
    assert direct in (None, 'prev', 'next')
    if order_types is not None:
      order_types_str = ','.join(t.value for t in order_types)
    else:
      order_types_str = None

    params = {
        "symbol": symbol,
        "types": order_types_str,
        "start-date": start_date,
        "end-date": end_date,
        "from": from_id,
        "direct": direct,
        "size": size
    }
    update = self._query(method, url, params)
    return update

  def query_single_order_match_results(self, order_id):
    method = 'GET'
    url = '/v1/order/orders/{order_id}/matchresults'.format(order_id=order_id)
    update = self._query(method, url)
    return update

  def place_order(
      self,
      *,
      account_id,
      symbol,
      price,
      amount,
      huobi_order_type,
  ):
    method = 'POST'
    url = '/v1/order/orders/place'
    params = {
        'account-id': account_id,
        'amount': amount,
        'price': price,
        'symbol': symbol,
        'type': huobi_order_type.value,
    }
    update = self._query(method, url, params)
    return update

  def cancel_order(self, order_id):
    method = 'POST'
    url = '/v1/order/orders/{order_id}/submitcancel'.format(order_id=order_id)
    update = self._query(method, url)
    return update

  def batch_cancel_orders(self, order_id_list: typing.List[str]):
    method = 'POST'
    url = '/v1/order/orders/batchcancel'
    assert len(order_id_list) <= 50, (len(order_id_list), order_id_list)

    params = {
        'order-ids': order_id_list,
    }
    update = self._query(method, url, params)
    return update

  def batch_cancel_open_orders(
      self,
      account_id=None,
      symbol=None,
      side=None,
      size=None,
  ):
    method = 'POST'
    url = '/v1/order/orders/batchCancelOpenOrders'
    assert bool(account_id) == bool(symbol), \
        'account-id and symbol must be both None or both set'
    if size is not None:
      assert 0 <= size <= 100, size
    side = side if side is None else side.value
    params = {
        'account-id': account_id,
        'symbol': symbol,
        'side': side,
        'size': size,
    }
    update = self._query(method, url, params)
    return update


if __name__ == "__main__":
  client = HuobiNativePrivateClient(
      key_file='../../coin_key/view/prestoinv0/Huobi/view_key.json')
  update = client.query_accounts()
  account_ids = [data['id'] for data in update.msg['data'] if data['type'] == 'super-margin']
  print(client.query_account_balance(account_ids[0]))
  account_ids = [data['id'] for data in update.msg['data'] if data['type'] == 'margin']
  print(client.query_account_balance(account_ids[0]))
